#/*************************************************************************
#    > File Name: calc/calc_setting.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: 2022年07月21日 星期四 14时31分39秒
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import paramiko 
sys.path.append('../Lib')
from Load_Lib import * 
from event_path import * 

# 设置print的显示规则
np.set_printoptions(threshold=np.inf)
pd.set_option('display.max_columns',None)
pd.set_option('display.max_rows',None)

# 设置基本的工作路径
#work_cluster="college"
work_cluster="yancy"
#mg_event_method = "pbs"
mg_event_method = "slurm"
#mg_event_method = "seperate"
#mg_event_method = "scan"

if work_cluster == "college":
	work_dir="/public/home/yancy/Code/Git/auto_mab/Lib/calc/"
	MG_dir="/public/home/yancy/Software/MG5/"
	ARCH_dir="/public/home/yancy/Software/ARCH/"
	BASDA_dir="/public/home/yancy/Software/BASDA/"
	node_name=["node1","node2", "node3","node4","node5","node6","node7","node8","node9","node10","node11","node12"]
elif work_cluster == "yancy":
	work_dir= "/home/yancy/Software/ARCH_test/Auto_Running/auto_mab/Lib/calc/"
	MG_dir=   "/home/yancy/Software/MG5/"
	ARCH_dir= "/home/yancy/Code/ARCH_test/"
	BASDA_dir="/home/yancy/Code/BASDA/"
	node_name=["node1", "node2","gpu"]
node_num=len(node_name) 



# 设置基本参数

process="THDM_CH_4b0j"
param_kind="mh_mch"
#param_kind="sba_tb"
parton_pt="20"
#param_mode="fixed_width"
param_mode="seperate"


# 组合名称
event_kind=param_kind+"_"+parton_pt
process_event_kind=process+"_"+event_kind


# 设置MG param card的参数
process_input_folder=work_dir+'input_{fprocess}/'.format(fprocess=process)
process_output_folder=work_dir+'output_{fprocess}/'.format(fprocess=process)
param_source = process_input_folder+'/card/param_card.dat'
if param_mode=="fixed_width":
	MGparam_label_value={"Block frblock":2,"23 9.118760e+01 # MZ ":4}
	MGparam_replace_value={"tb" :"    1 scan1: {} # TB",
						  "sba" :"    2 scan1: {} # sinbma",
						  "mh"  :"   25 scan1: {} # mh1",
						  "mH"  :"   35 scan1: {} # mh2",
						  "mA"  :"   36 scan1: {} # mh3",
						  "mch" :"   37 scan1: {} # mhc"}
	MGparam_label_constant={"DECAY  24 2.085000e+00 # WW ":4}
	MGparam_replace_constant={"wh"  :"DECAY  25 Auto # Wh1 ",
							  "wH"  :"DECAY  35 Auto # Wh1 ",
							  "wA"  :"DECAY  36 Auto # Wh1 ",
							  "wch" :"DECAY  37 Auto# Wh1 "}
elif param_mode=="variable_width":
	MGparam_label_value={"Block frblock":2,"23 9.118760e+01 # MZ ":4,"DECAY  24 2.085000e+00 # WW ":4}
	MGparam_replace_value={"tb" :"    1 scan1: {} # TB",
						  "sba" :"    2 scan1: {} # sinbma",
						  "mh"  :"   25 scan1: {} # mh1",
						  "mH"  :"   35 scan1: {} # mh2",
						  "mA"  :"   36 scan1: {} # mh3",
						  "mch" :"   37 scan1: {} # mhc",
						  "wh"  :"DECAY  25 Auto # Wh ",
						  "wH"  :"DECAY  35 Auto # WH ",
						  "wA"  :"DECAY  36 Auto # WA ",
						  "wch" :"DECAY  37 Auto # Wch "}
elif param_mode=="seperate":
	MGparam_label_value={"Block frblock":2,"23 9.118760e+01 # MZ ":4}
	MGparam_replace_value={"tb":"      1 {} # TB",
			              "sba" :"    2 {} # sinbma",
				          "mh"  :"   25 {} # mh1",
					      "mH"  :"   35 {} # mh2",
						  "mA"  :"   36 {} # mh3",
						  "mch" :"   37 {} # mhc"}
	MGparam_label_constant={"DECAY  24 2.085000e+00 # WW ":4}
	MGparam_replace_constant={"wh"  :"DECAY  25 Auto # Wh1 ",
							  "wH"  :"DECAY  35 Auto # Wh1 ",
							  "wA"  :"DECAY  36 Auto # Wh1 ",
							  "wch" :"DECAY  37 Auto# Wh1 "}
scan_folder="scan_param_{fname}_{fpt}/".format(fname=param_kind, fpt=parton_pt)
param_general_file_name=process_output_folder+scan_folder+'/param_card_{}.dat'




# 设置parton level 截面计算参数
theoretical_para_file_name=process_input_folder+"/data/data_grid_{fname}.tsv".format(fname=param_kind)
param_df=Load_THDM_Param(theoretical_para_file_name)
param_length=len(param_df)


if(param_kind=="mh_mch"):
	df_para=["mh","mch"]
	xaxis=param_df["mh"].to_list()
	yaxis=param_df["mch"].to_list()
	xname="M_{h}"
	yname="M_{H^{#pm}}"
	bins=[10,20,120,9,80,170]
elif(param_kind=="sba_tb"):
	df_para=["sba","tb"]
	xaxis=param_df["sba"].to_list()
	yaxis=param_df["tb"].to_list()
	xname="sin(#beta-#alpha)"
	yname="tan(#beta)"
	bins=[10,-0.3,-0.0,12,0,60]

if parton_pt=="10":
	MG_run_card=process_input_folder+"/card/run_card_10.dat"
	MG_delphes_card=process_input_folder+"/card/delphes_card_10.dat"
elif parton_pt=="20":
	MG_run_card=process_input_folder+"/card/run_card_20.dat"
	MG_delphes_card=process_input_folder+"/card/delphes_card_20.dat"




# 设置MG工作路径
event_folder="Events_{fname}_{fpt}".format(fname=param_kind,fpt=parton_pt)
MG_parallel_output_folder=process_output_folder+"parallel_tmp/mg_results_{fname}_{fpt}/".format(fname=param_kind, fpt=parton_pt)


# 对于考虑 cut 的事例
## 设置集群节点参数，用于并行计算
if mg_event_method == "scan":
	folder_num=node_num
elif mg_event_method == "seperate":
	folder_num=param_length
elif mg_event_method == "slurm":
	folder_num=param_length
elif mg_event_method == "pbs":
	folder_num=param_length
else:
	print("please add correct mg_event_method, and rerun the code. The current value is ", mg_event_method)
	exit(0)

MG_original_folder=MG_dir+"{fprocess}".format(fprocess=process)
MG_split_event_folder=[]
MG_combined_event_folder=MG_dir+event_folder+"/{fprocess}_0/Events/".format(fprocess=process)
for num in range(folder_num):
	new_folder=MG_dir+event_folder+"/{fprocess}_{fn}/".format(fprocess=process,fn=num+1)
	MG_split_event_folder.append(new_folder)


# 对于考虑 no cut 的事例
event_folder_nocut="Events_{fname}_nocut".format(fname=param_kind)
MG_original_folder_nocut=MG_dir+"/{fprocess}_nocut/Events/".format(fprocess=process)
MG_combined_event_folder_nocut=MG_dir+event_folder_nocut+"/{fprocess}_0/Events/".format(fprocess=process)
MG_split_event_folder_nocut=[]
for num in range(folder_num):
	new_folder=MG_dir+event_folder_nocut+"/{fprocess}_nocut_{fn}/".format(fprocess=process,fn=num+1)
	MG_split_event_folder_nocut.append(new_folder)


scan_general_file_name="scan_run_[01-{fnum}]_mycode.txt"

# 设置parton level截面与其他输入参数作为hotmap的参数。 
MG_hotmap_out_file=process_output_folder+"/xection/{fname}_for_MG_xection.dat".format(fname=process_event_kind)
MG_hotmap_variable_name="Parton_xection"
MG_hotmap_file_name=process_output_folder+"/hotmap/{fname}_for_MG_xection.C".format(fname=process_event_kind)
MG_param_out_file=process_output_folder+"/param/{fname}_for_MG.dat".format(fname=process_event_kind)

MG_para_plot={	"file_name": MG_hotmap_file_name,
		"bins":bins,
		"data_length":param_length,
		"var_name":MG_hotmap_variable_name,
		"X_name": xname,
		"Y_name": yname,
		}

# 设置ARCH相关参数
ARCH_control_template= './examples/test_input/control_ARCH'
ARCH_target_control_folder= ARCH_dir+'/control/control_scan_{}'.format(process_event_kind)
ARCH_event_general_folder="events/Event_Analyse_scan_{fname}/"
ARCH_control_dir={"source":ARCH_control_template,"target":ARCH_target_control_folder}
ARCH_para_path={"origin": {"process":"PROCESS          : {}_$fnum".format(process_event_kind),
			"out_folder" :"OUTPUT_FOLDER    : "+ARCH_event_general_folder.format(fname=process_event_kind),
			"tmp" :"RECORD_FILE      : 'tmp/{}_$fnum.dat'".format(process_event_kind),
			"in_folder" :"INPUT_FOLDER     : './control/control_scan_{}'".format(process_event_kind), #这里要和上面定义的ARCH_target_control_folder对应
			"file":"INPUT_FILE       : 'file_$fnum.dat'"},
		"source":"path.dat",
		"labels":{"PROCESS          : process":0,
			"OUTPUT_FOLDER    : 'events/Event_Analyse_scan_process'":0,
			"RECORD_FILE      : 'tmp/sig_process.dat'":0,
			"INPUT_FOLDER     : './control/control_scan_process'":0,
			"INPUT_FILE       : 'file_process.dat'":0}
			}
ARCH_para_file={"origin": {"file":"FILE            : [\""+MG_combined_event_folder+"run_$fnum/tag_1_delphes_events.root\"]",
			"cross":"CROSS_SECTION   : $fvalue"},
			"source":"file.dat",
			"labels":{"FILE            : ['sig_Wh_4b_1_1.root']":0,
			"CROSS_SECTION   : 32.0566":0}
			}
ARCH_replaced_para={"path":ARCH_para_path, 
		"file":ARCH_para_file}


ARCH_cut_file_name  =MG_combined_event_folder      +scan_general_file_name.format(fnum=param_length)
ARCH_nocut_file_name=MG_combined_event_folder_nocut+scan_general_file_name.format(fnum=param_length)
ARCH_parallel_output_folder=process_output_folder+"parallel_tmp/ARCH_results_{fname}_{fpt}/".format(fname=param_kind, fpt=parton_pt)


# 设置BASDA相关参数，利用ARCH输出结果

ARCH_tmp_folder= ARCH_dir+ARCH_event_general_folder.format(fname=process_event_kind)
ARCH_tmp_name = ARCH_tmp_folder+'/tmp/sigh_w4b_{}.dat'
ARCH_summary_folder= process_output_folder+"/BASDA_para_{event}".format(event=process_event_kind)

ARCH_hotmap_out_file=process_output_folder+"/xection/{fname}_for_ARCH_xection.dat".format(fname=process_event_kind)
ARCH_hotmap_variable_name="ARCH_xection"
ARCH_hotmap_file_name=process_output_folder+"/hotmap/{fname}_for_ARCH_xection.C".format(fname=process_event_kind)
ARCH_param_out_file=process_output_folder+"/param/{fname}_for_ARCH.dat".format(fname=process_event_kind)


ARCH_para_plot={"file_name": ARCH_hotmap_file_name,
		"bins":bins,
		"data_length":param_length,
		"var_name":ARCH_hotmap_variable_name,
		"X_name": xname,
		"Y_name": yname,
		}






BASDA_control_template= './examples/test_input/control_BASDA'
BASDA_target_control_folder= BASDA_dir+'/control/control_scan_{}'.format(process_event_kind)
BASDA_event_general_folder="events/Event_Analyse_scan_{fname}"
BASDA_control_dir={"source":BASDA_control_template,"target":BASDA_target_control_folder}

BASDA_para_path={"origin": {"process":"PROCESS                   : {}_$fnum".format(process_event_kind),
			"out_folder" :"OUTPUT_FOLDER    : "+BASDA_event_general_folder.format(fname=process_event_kind)+"_$fnum",
			    "tmp" :"RECORD_FILE               : \"tmp/{}_$fnum.dat\"".format(process_event_kind),
			"in_folder" :"INPUT_FOLDER     : './control/control_scan_{}'".format(process_event_kind), #这里要和上面定义的ARCH_target_control_folder对应
			    "file":"INPUT_FILE                : \"file_$fnum.dat\""},
		"source":"path.dat",
		"labels":{"PROCESS                   : 10_bp3_4b0j_1":0,
                    "OUTPUT_FOLDER             : 'events/Event_scan_1'":0,
                    "RECORD_FILE               : 'tmp/tmp_h_w4b_1.dat'":0, 
		    "INPUT_FOLDER              : './control/control_scan_4b0j'":0,
                    "INPUT_FILE                : 'file_1.dat'":0}
			}

BASDA_para_file={"origin": {"file":"FILE_0            : xsection_$fnum.dat"},
			"source":"file.dat",
			"labels":{"FILE_0            : xsection_1.dat":0}
			}
BASDA_para_xection={"origin": {"file":"        "+ARCH_tmp_folder+"events/sigh_w4b_$fnum.root        : $fvalue "},
			"source":"xsection.dat",
			"labels":{"        file  : 1.554":0}
			}


BASDA_replaced_para={"path":BASDA_para_path, 
		"file":BASDA_para_file,
		"xection":BASDA_para_xection}

# 提取最终significance，利用BASDA的输出结果
BASDA_out_name = BASDA_dir+BASDA_event_general_folder.format(fname=process_event_kind)+"_{}"+"/data/125/output_table_combined_125.dat"
BASDA_hotmap_out_file=process_output_folder+"/xection/{fname}_for_BASDA_significance.dat".format(fname=process_event_kind)
BASDA_hotmap_variable_name="sig_data"
BASDA_hotmap_file_name=process_output_folder+"/hotmap/{fname}_for_BASDA_significance.C".format(fname=process_event_kind)
BASDA_param_out_file=process_output_folder+"/param/{fname}_for_BASDA.dat".format(fname=process_event_kind)


BASDA_parallel_output_folder=process_output_folder+"parallel_tmp/BASDA_results_{fname}_{fpt}/".format(fname=param_kind, fpt=parton_pt)

# 总结默认MG参数
MG_para={}
MG_para["MG_dir"]=MG_dir
MG_para["work_dir"]=work_dir
MG_para["ARCH_dir"]=ARCH_dir
MG_para["BASDA_dir"]=BASDA_dir
MG_para["replace_value"]=MGparam_replace_value
MG_para["data_df"]=param_df
MG_para["label_value"]=MGparam_label_value
MG_para["node_num"]=node_num
MG_para["folder_num"]=folder_num
MG_para["node_name"]=node_name
MG_para["length"]=param_length
MG_para["source"]=param_source
MG_para["general_file_name"]=param_general_file_name
MG_para["run_card"]=MG_run_card
MG_para["delphes_card"]=MG_delphes_card
MG_para["record_file"]=MG_parallel_output_folder+"mg_{}"
MG_para["scan_file_general_name"]=scan_general_file_name

MG_para["original_folder"]=MG_original_folder
MG_para["split_event_folder"]=MG_split_event_folder
MG_para["combined_event_folder"]=MG_combined_event_folder
MG_para["param_out_file"]=MG_param_out_file
MG_para["combined_scan_file"]=MG_para["combined_event_folder"]+scan_general_file_name
MG_para["split_scan_file"]=[this_file+scan_general_file_name for this_file in MG_split_event_folder]

MG_para["original_folder_nocut"]=MG_original_folder_nocut
MG_para["split_event_folder_nocut"]=MG_split_event_folder_nocut
MG_para["combined_event_folder_nocut"]=MG_combined_event_folder_nocut
MG_para["combined_scan_file_nocut"]=MG_para["combined_event_folder_nocut"]+scan_general_file_name
MG_para["split_scan_file_nocut"]=[this_file+scan_general_file_name for this_file in MG_split_event_folder_nocut]
MG_para["param_out_file_nocut"]=MG_param_out_file+"_nocut"
